import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import paddle
from paddle.vision.transforms import Compose, Normalize
from sklearn.metrics import accuracy_score  # 仅用于模型验证示例

plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定黑体（SimHei）
plt.rcParams['axes.unicode_minus'] = False  # 解决负号乱码
plt.rcParams["font.family"] = ["SimHei"]

# ---------------------- 1. 数据加载与基础清洗 ----------------------
def load_and_clean_image(img_path, target_size=(224, 224), denoise_method="median"):
    """
    加载图像并进行基础清洗（去噪、尺寸归一化）
    :param img_path: 图像路径
    :param target_size: 目标尺寸（H, W）
    :param denoise_method: 去噪方法（median/gaussian）
    :return: 清洗后的图像（numpy数组，RGB格式）
    """
    try:
        # 读取图像（支持中文路径）
        img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError(f"无法加载图像: {img_path}")
        
        # BGR转RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # 去噪处理
        if denoise_method == "median":
            img = cv2.medianBlur(img, 3)  # 中值滤波去椒盐噪声
        elif denoise_method == "gaussian":
            img = cv2.GaussianBlur(img, (3, 3), 0)  # 高斯滤波去高斯噪声
        
        # 尺寸归一化
        img = cv2.resize(img, target_size)
        
        return img
    except Exception as e:
        print(f"处理图像失败: {e}")
        return None

# ---------------------- 2. 数据增强（使用Paddle视觉工具） ----------------------
def build_augmentation():
    """定义数据增强流程（训练集专用）"""
    transform = Compose([
        paddle.vision.transforms.RandomRotation(degrees=10),        # 随机旋转±10度
        paddle.vision.transforms.RandomHorizontalFlip(prob=0.5),    # 水平翻转（概率0.5）
        paddle.vision.transforms.ColorJitter(brightness=0.1),      # 亮度调整
        paddle.vision.transforms.ToTensor(),                       # 转为Tensor（HWC→CHW）
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 标准化
    ])
    return transform

# ---------------------- 3. 可视化对比函数 ----------------------
def visualize_comparison(original_img, cleaned_img, augmented_img=None, title="清洗效果对比"):
    """可视化原始图像、清洗后图像、增强后图像对比"""
    plt.figure(figsize=(18, 6))
    
    # 原始图像
    plt.subplot(1, 3, 1)
    plt.imshow(original_img)
    plt.title("原始图像")
    plt.axis("off")
    
    # 清洗后图像
    plt.subplot(1, 3, 2)
    plt.imshow(cleaned_img)
    plt.title("清洗后图像（去噪+Resize）")
    plt.axis("off")
    
    # 增强后图像（可选）
    if augmented_img is not None:
        # 转换Tensor为numpy并反标准化
        augmented_img = augmented_img.transpose((1, 2, 0)).numpy()
        augmented_img = augmented_img * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406])
        augmented_img = np.clip(augmented_img, 0, 1)
        
        plt.subplot(1, 3, 3)
        plt.imshow(augmented_img)
        plt.title("数据增强后图像")
        plt.axis("off")
    
    plt.suptitle(title, fontsize=16)
    plt.tight_layout()
    plt.show()

# ---------------------- 4. 图像质量评估 ----------------------
def calculate_psnr_ssim(original_img, cleaned_img, target_size=(224, 224)):
    """计算PSNR和SSIM指标（确保尺寸一致）"""
    # 调整原始图像尺寸
    original_img_resized = cv2.resize(original_img, target_size)
    
    # 转换为灰度图
    original_gray = cv2.cvtColor(original_img_resized, cv2.COLOR_RGB2GRAY)
    cleaned_gray = cv2.cvtColor(cleaned_img, cv2.COLOR_RGB2GRAY)
    
    # PSNR
    psnr = cv2.PSNR(original_gray, cleaned_gray)
    
    # 基于直方图的相似度（替代SSIM）
    hist1 = cv2.calcHist([original_gray], [0], None, [256], [0, 256])
    hist2 = cv2.calcHist([cleaned_gray], [0], None, [256], [0, 256])
    ssim = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
    
    return psnr, ssim

# ---------------------- 5. 模型验证（分类任务示例） ----------------------
def evaluate_model(cleaned_imgs, labels, model=None):
    """简单模型验证清洗效果（此处为示例，需替换为实际模型）"""
    # 模拟模型：计算清洗后图像的平均亮度作为分类特征
    features = np.array([img.mean() for img in cleaned_imgs])
    pred_labels = (features > 128).astype(int)  # 假设亮度高为无缺陷
    acc = accuracy_score(labels, pred_labels)
    return acc

# ---------------------- 6. 主流程运行 ----------------------
if __name__ == "__main__":
    # 模拟数据：假设3张含噪声的PCB图像路径
    img_paths = ["examples/cat9.png", "examples/cat9.png" ]
    labels = [1,1]  # 1:无缺陷，0:有缺陷
    
    # 构建数据集（清洗+增强）
    transform = build_augmentation()
    cleaned_imgs = []
    augmented_imgs = []
    
    for path in img_paths:
        # 加载并清洗图像
        cleaned_img = load_and_clean_image(path)
        if cleaned_img is None:
            continue
        cleaned_imgs.append(cleaned_img)
        
        # 应用数据增强（转为Tensor后处理） 

        img_tensor = paddle.to_tensor(cleaned_img.transpose((2, 0, 1)).astype(np.float32) / 255.0)# HWC→CHW并归一化
        augmented_tensor = transform(img_tensor)
        augmented_imgs.append(augmented_tensor)
    
    # 可视化第一张图像的清洗效果
    original_img = cv2.cvtColor(cv2.imread(img_paths[0]), cv2.COLOR_BGR2RGB)
    visualize_comparison(original_img, cleaned_imgs[0], augmented_imgs[0])
    
    # 打印质量指标
    psnr, ssim = calculate_psnr_ssim(original_img, cleaned_imgs[0])
    print(f"PSNR: {psnr:.2f} dB, SSIM: {ssim:.2f}")
    
    # 模型验证
    acc = evaluate_model(cleaned_imgs, labels)
    print(f"清洗后模型准确率: {acc:.4f}")